package com.zh.stargraph2.core.job;

import cn.hutool.core.date.DateTime;
import com.zh.stargraph2.comfyui.client.api.ComfyuiApi;
import com.zh.stargraph2.core.pojo.ComfyuiTask;
import com.zh.stargraph2.core.service.RedisService;
import lombok.extern.log4j.Log4j2;
import org.redisson.api.RLock;
import org.redisson.api.RedissonClient;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Component;
import retrofit2.Call;
import retrofit2.Response;

import java.util.HashMap;

@Component
@Log4j2
public class RunTaskJob {

    @Autowired
    RedisService redisService;
    @Autowired
    ComfyuiApi comfyuiApi;
    @Autowired
    RedissonClient redissonClient;

    public final static String SPRING_TASK_LOCK_KEY = "SPRING_TASK_LOCK_KEY";

    private void addTaskToQueue() {
        try {
            ComfyuiTask task = redisService.popQueueTask();
            if (task == null) {
                return;
            }
            Call<HashMap> addQueueTask = comfyuiApi.addQueueTask(task.getComfyuiRequestDto());
            Response<HashMap> response = addQueueTask.execute();
            if (response.isSuccessful()) {
                HashMap data = response.body();
                task.setPromptId(data.get("prompt_id").toString());
                log.info("添加任务到Comfyui成功：{}", task.getPromptId());
                redisService.addStartedTask(task.getPromptId(), task);
            } else {
                String error = response.errorBody().string();
                log.error("添加任务到Comfyui错误：{}", error);
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    @Scheduled(cron = "*/1 * * * * ?")
    public void task() {
        RLock lock = redissonClient.getLock(SPRING_TASK_LOCK_KEY);
        if (lock.tryLock()) {
            try {
                addTaskToQueue();
                //System.out.println(DateTime.now());
            } finally {
                lock.unlock();
            }
        }
    }
}
